package sub

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"quipus/cmd/qpc/sub"
	"quipus/internal/utils"
	v1pb "quipus/proto/gen/api/v1"
	cpb "quipus/proto/gen/common"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/ollama/ollama/readline"
	"google.golang.org/protobuf/proto"

	"github.com/spf13/cobra"
)

var (
	subjectId   int32
	topicId     int32
	content     string
	interactive string
	url         string
	mode        string
)

func init() {
	chatCmd.PersistentFlags().Int32VarP(&subjectId, "subject", "s", 0, "the subject id of subject")
	chatCmd.PersistentFlags().Int32VarP(&topicId, "topic", "t", 0, "the topic id of topic")
	chatCmd.PersistentFlags().StringVarP(&content, "content", "c", "", "the chat content")
	chatCmd.PersistentFlags().StringVarP(&interactive, "interact-mode", "i", "NONE", "interactive mode: chat query none")
	chatCmd.PersistentFlags().StringVarP(&mode, "mode", "m", "NAIVE", "rag mode")

	knotCmd.AddCommand(chatCmd)
}

var chatCmd = &cobra.Command{
	Use:   "run",
	Short: "Run the cmd to geneate knot",
	RunE: func(cmd *cobra.Command, args []string) error {
		var chatMode cpb.RagMode
		if mode != "" {
			chatMode = cpb.RagMode(cpb.RagMode_value[strings.ToUpper(mode)])
		}
		qconn, e := sub.NewGrpcConnect(sub.QpConfig.QPAddress)
		if e != nil {
			return e
		}

		qclient := v1pb.NewKnotServiceClient(qconn)
		ctxWithToken := sub.NewCtxWithToken(sub.QpConfig.Token)
		subjectUrn := ""
		if subjectId > 0 {
			subjectUrn = utils.BuildUrn(utils.SubjectUrnPrefix, subjectId)
		}
		topicUrn := ""
		if topicId > 0 {
			topicUrn = utils.BuildUrn(utils.TopicUrnPrefix, topicId)
		}
		in := &v1pb.CreateKnotRequest{SubjectUrn: subjectUrn, TopicUrn: topicUrn, Content: content}
		im := strings.ToUpper(interactive)
		switch im {
		case "CHAT":
			in.InteractType = cpb.InteractType_CHAT
		case "QUERY":
			in.InteractType = cpb.InteractType_QUERY
		case "NONE":
			in.InteractType = cpb.InteractType_INTERACT_UNSPECIFIED
		}

		rsp, err := qclient.CreateKnot(ctxWithToken, in)
		if err != nil {
			return err
		}
		if in.InteractType != cpb.InteractType_CHAT {
			fmt.Print(rsp)
		} else {
			fmt.Println("interactive mode is required")
			fmt.Printf(">>> %s\n", content)
			RunInteractive(rsp, chatMode)
		}

		return nil
	},
}

func RunInteractive(knot *v1pb.Knot, chatMode cpb.RagMode) {
	header := http.Header{}
	header.Add("Authorization", fmt.Sprintf("Bearer %s", sub.QpConfig.Token))
	c, _, err := websocket.DefaultDialer.Dial(sub.QpConfig.WS+"/"+knot.Uid, header)
	if err != nil {
		log.Fatal("error:", err)
	}

	SendViaWebSocket(c, "", knot.Urn, v1pb.Phase_HELLO, chatMode)
	ctx, cancel := context.WithCancel(context.Background())
	ReadWebSocket(ctx, c)

	scanner, err := readline.New(readline.Prompt{
		Prompt:         ">>> ",
		AltPrompt:      "... ",
		Placeholder:    "Send a message (/q for quit)",
		AltPlaceholder: `Use """ to end multi-line input`,
	})

	defer cancel()

	if err != nil {
		log.Fatal("error:", err)
	}

	for {
		line, err := scanner.Readline()
		if err != nil {
			break
		}

		if strings.Trim(line, " ") == "" {
			continue
		}

		if line == "/q" {
			SendViaWebSocket(c, "", knot.Urn, v1pb.Phase_Close, chatMode)
			c.Close()
			break
		}

		if line == "/c" {
			ReadWebSocket(ctx, c)
			continue
		}

		SendViaWebSocket(c, line, knot.Urn, v1pb.Phase_Continue, chatMode)
		ReadWebSocket(ctx, c)
	}
}
func SendViaWebSocket(conn *websocket.Conn, content string, urn string, phase v1pb.Phase, chatMode cpb.RagMode) {
	chatMsg := v1pb.Exchange{Content: content, AskId: urn, Phase: phase, Mode: chatMode}

	b, e := proto.Marshal(&chatMsg)
	if e != nil {
		log.Fatal("error:", e)
	}
	conn.WriteMessage(websocket.TextMessage, b)
}

func ReadWebSocket(ctx context.Context, conn *websocket.Conn) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				break
			}
			chatMsg := v1pb.Exchange{}

			err = proto.Unmarshal(message, &chatMsg)
			if err != nil {
				log.Println("unmarshal:", err)
				break
			}
			if chatMsg.Phase == v1pb.Phase_End {
				fmt.Println(chatMsg.Content)
				return
			}
			fmt.Print(chatMsg.Content)
		}
	}
}
